% This program calculates the moments in the model we attempt to match to
% the data
function dataset = mechanical_effect(Prob, Ss, rs, lambdas, tol, max_It);
    
    global bet0 bet1 a b w e pi_theta N_a N_w N_e w_max w_min alpha...
        pi_e Pr skill_shock r_base xi1 myiter...
        shares w_bar ss_calc unemp zeta myslope elasts nu xi eta mu_a thetas tau sigmae share log_e_sigma pe
    
    eta = 1;
    ln_e_hat = -1/2*log_e_sigma^2; % mean of the non-normalized idiosyncratic wage shock
    N_n = length(Ss);
    
    % Dataset for the various moments
    dataset = cell(5, 23);

    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % Steady-state choices and distribution
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    w_bar = dot(w, sum(Prob,2));
    
    [P_S, F_es, wprimeS, educ_mat] = choiceProb(Ss, rs);
    % Steady-state choice of neighborhood
    P_S0 = P_S;
    educ_mat0 = educ_mat;
    % Initial Distribution
    myDist0 = P_S .* Prob;
    myDist0  = sum(bsxfun(@times, myDist0 , reshape(pi_theta, [1,1,1,2])),4);
    F_es0 = F_es;
    w0 = w;
    
    denom = sum(myDist0 .* wprimeS, 1:3);
    bet1 = bet1 / (denom)^xi1;
    Ss = Ss .* (denom)^xi1;

    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % Invert the housing equation to get lambdas
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    lambdas =  inv_hs(shares, rs);
   
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % Calculate the wage cutoff associated with college
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    edu_vec = educ_mat(:);
    pi_joint_vec = myDist0(:);
    
    edu_share_col = 0.178;
    [edu_dec, index_edu] = sort(edu_vec,'descend');
    pi_joint_dec = pi_joint_vec(index_edu); %probability of sorted education

    cdf_edu = cumsum(pi_joint_dec);

    [unique_edu_dec,index,~] = unique(edu_dec,'first','legacy');
    find_edu_college = @(x)interp1(unique_edu_dec,cdf_edu(index),x,[],'extrap')-edu_share_col;
    edu_college = fzero(find_edu_college,median(edu_dec)); 
      
    for t = 1:5
        myiter = t;
        
        % Generate shock to spillover
        if t == 2
            eta = skill_shock;    
            F_es0 = F_es0 * skill_shock; % Adjust parent's wage for shock
        end
        
        pi_w = sum(Prob, 2);
        w_bar = dot(w, pi_w);
        
        w0_1 = w0;
        w0 = w;
        
        w_max0 = w_max;
        w_min0 = w_min;       
        if t == 1
            [Ss, rs] = find_spillovers_dyn(Ss, lambdas, rs, Prob, tol, max_It);
            [~, F_es, wprimeS, educ_mat] = choiceProb(Ss, rs);
        else
            
            for iter = 1:40
                [~, F_es, wprimeS, ~] = choiceProb_no_ed(Ss, rs, educ_mat0);
                myDist0 = P_S0 .* Prob;
                wMass = squeeze(sum(sum(bsxfun(@times, wprimeS, myDist0), 1:2), 4));
                popMass = squeeze(sum(sum(myDist0, 1:2),4));
                Ssi = wMass ./ popMass;
                Ssi = Ssi(:).^xi1; 
                d_S = abs(Ssi - Ss);
                if sum(d_S) < 2 * tol
                    break
                end
                Ss = Ssi;
            end
            P_S = P_S0;
            [~, F_es, wprimeS, educ_mat] = choiceProb_no_ed(Ss, rs, educ_mat0);
            P_S = P_S0;
            
        end
        
        N_size = squeeze(sum(bsxfun(@times, Prob .* P_S , reshape(pi_theta, [1,1,1,2])),[1,2,4]));
        
        % Adjust parent's wage for the skill shock but don't let it impact
        % their optimal decisions
        if t == 1
            Prob_kids = update_dist_int_edu_dn_Ss(Prob, P_S0, educ_mat0, Pr, Ss, rs, 1);
        else
            Prob_kids = update_dist_int_edu_dn_Ss(Prob, P_S0, educ_mat0, Pr, Ss, rs, 0);
        end
        
        % Calculate the marginal distribution of w
        pi_w = sum(Prob, 2);
        cdf_w = cumsum(pi_w,1);
        
        % Find the median
        find_w_median= @(x)interp1(w0,cdf_w,x,[],'extrap')-.5;
        w_median = fzero(find_w_median,(w_max+w_min)/2);
        
        b0 = b;
        
        w_m = dot(w,pi_w);
        w_v = dot(((w-w_m).^2),pi_w);
        sd_w = w_v^(1/2);
        cofv = sd_w / w_m;
        
        % Distribution by neighborhood
        myDist1 = P_S .* Prob;
        myDist1  = sum(bsxfun(@times, myDist1 , reshape(pi_theta, [1,1,1,2])),4);
        
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % Return to Education
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        share_C = sum((educ_mat >= edu_college) .* myDist1,1:2)./sum( myDist1,1:2);
        Elogw_C = (sum((educ_mat > edu_college) .* log(wprimeS) .* myDist1,1:3)/sum((educ_mat > edu_college) .* myDist1,1:3));
        Elogw_NC = (sum((educ_mat < edu_college) .* log(wprimeS) .* myDist1,1:3)/sum((educ_mat < edu_college) .* myDist1,1:3));
        if ~isnan(edu_college)
            total_educ_return = Elogw_C-Elogw_NC;
        else
            total_educ_return=nan;
        end
        
        col_shares = squeeze(sum((educ_mat >= edu_college) .* myDist1,1:2)./sum(myDist1,1:2))';
        
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % Average Education and ability
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%         
        
        avg_ed = squeeze(sum(educ_mat.* myDist1,1:2)./sum( myDist1 ,1:2))';
        avg_ability = squeeze(sum(a'.* myDist1,1:2)./sum( myDist1 ,1:2))';
             
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % Gini Index
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        my_prod = pi_w .* pi_w' .* abs(w0 - w0');
        gini = sum(my_prod(:))/(2*dot(pi_w, w0));
        
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % Income Dissimilarity Index
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        rich_cutoff = 0.8;
        poor = rich_cutoff; 
        rich = 1 - rich_cutoff; 
        agg_cdf = cumsum(myDist1, 1);
        
        P_poor_S = zeros(length(Ss),1);
        P_rich_S = zeros(length(Ss),1);
        
        find_rich_w = @(x)interp1(w0,cdf_w,x,[],'extrap')-rich_cutoff;
        w_rich = fzero(find_rich_w,(w_max0+w_min0)/2);
        
        share_h = squeeze(sum((w0 < w_rich) .* myDist1,1:2))/ sum(squeeze(sum((w0 < w_rich) .* myDist1,1:2)));
        share_l = (squeeze(sum((w0 >= w_rich) .* myDist1,1:2)))/(sum(squeeze(sum((w0 >= w_rich) .* myDist1,1:2))));
        
        dissim = sum(abs(share_h - share_l))*0.5; 
        
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % Dissimilarity Index by Education
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        share_h = squeeze(sum((educ_mat < edu_college) .* myDist1,1:2))/ sum(squeeze(sum((educ_mat < edu_college) .* myDist1,1:2)));
        share_l = (squeeze(sum((educ_mat >= edu_college) .* myDist1,1:2)))/(sum(squeeze(sum((educ_mat >= edu_college) .* myDist1,1:2))));
        
        dissim_edu = sum(abs(share_h - share_l))*0.5;
        
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % Rank-Rank Correlation
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        %% Upward Mobility
        % find 25th percentile (cdf_w was already defined above)
        [~,idx] = min(abs(cdf_w -.25));
        pi_w_kids = sum(Prob_kids,2);
        cdf_w_kids = cumsum(pi_w_kids);
        find_25_w_kids = @(x)interp1(w,cdf_w_kids,x,[],'extrap')-0.25;
        w_25_kids = fzero(find_25_w_kids,(w_max+w_min)/2);

        % find probability of staying below 25th percentile for each (ability,wage) pair
        Q1toQ = zeros(N_a,N_w, N_n);
        
        for t2=1:2
            for i=1:N_a
                for j=1:idx
                    for c = 1:N_n
                        my_share = Prob(j,i)*P_S(j,i,c,t2)*pi_theta(t2);
                        
                        wj = w0(j);
                        raw_Omega = unemp + (b0   + F_es(j,i,c))  .* f_w(wj) ;
                        Q1toQ(i,j) = Q1toQ(i,j) + ...
                            my_share*normcdf((log(w_25_kids) - log(raw_Omega) - ln_e_hat)/log_e_sigma);
                    end
                end
            end
        end
        mobility_up = sum(Q1toQ(:))/cdf_w(idx);
 

        % Find 25th and 75th percentile
        [~,find_w_p25]= min(abs(cdf_w-0.25));
        w_p25 = w0(find_w_p25);
        
        [~,find_w_p75]= min(abs(cdf_w-0.75));
        w_p75 = w0(find_w_p75);
        
        % Average wage of child with parent at 25th percentile
        % Interpolate between grid points above and below percentile
        idx = zeros(length(w0),1);
        idx(find_w_p25) = 1;
       
        p_blm_1 = sum((idx .* myDist1 .* wprimeS), 1:3) ./ sum((idx .* myDist1), 1:3);
        if cdf_w(find_w_p25) < 0.25
            idx = zeros(length(w0),1);
            idx(find_w_p25 + 1) = 1;
            p_blm_2 = sum((idx .* myDist1 .* wprimeS), 1:3) ./ sum((idx .* myDist1), 1:3);
            wfirst = (0.25 - cdf_w(find_w_p25))/(cdf_w(find_w_p25+1) -  cdf_w(find_w_p25));
            wlast = (cdf_w(find_w_p25+1) - 0.25)/(cdf_w(find_w_p25+1) -  cdf_w(find_w_p25));
        else
            idx = zeros(length(w0),1);
            idx(find_w_p25 - 1) = 1;
            p_blm_2 = sum((idx .* myDist1 .* wprimeS), 1:3) ./ sum((idx .* myDist1), 1:3);
            wfirst = (cdf_w(find_w_p25) - 0.25)/(cdf_w(find_w_p25) -  cdf_w(find_w_p25-1));
            wlast = (0.25 - cdf_w(find_w_p25-1))/(cdf_w(find_w_p25) -  cdf_w(find_w_p25-1));
        end
        p_blm = p_blm_1*wlast + p_blm_2 *wfirst;
        
        % Average wage of child with parent at 75th percentile
        % Interpolate between grid points above and below percentile
        idx = zeros(length(w0),1);
        idx(find_w_p75) = 1;
        p_abm_1 = sum((idx .* myDist1 .* wprimeS), 1:3) ./ sum((idx .* myDist1), 1:3);
        if cdf_w(find_w_p75) < 0.75
            idx = zeros(length(w0),1);
            idx(find_w_p75 + 1) = 1;
            p_abm_2 = sum((idx .* myDist1 .* wprimeS), 1:3) ./ sum((idx .* myDist1), 1:3);
            wfirst = (0.75 - cdf_w(find_w_p75))/(cdf_w(find_w_p75+1) -  cdf_w(find_w_p75));
            wlast = (cdf_w(find_w_p75+1) - 0.75)/(cdf_w(find_w_p75+1) -  cdf_w(find_w_p75));
        else
            idx = zeros(length(w0),1);
            idx(find_w_p75 - 1) = 1;
            p_abm_2 = sum((idx .* myDist1 .* wprimeS), 1:3) ./ sum((idx .* myDist1), 1:3);
            wfirst = (cdf_w(find_w_p75) - 0.75)/(cdf_w(find_w_p75) -  cdf_w(find_w_p75-1));
            wlast = (0.75 - cdf_w(find_w_p75-1))/(cdf_w(find_w_p75) -  cdf_w(find_w_p75-1));
        end
        p_abm = p_abm_1*wlast + p_abm_2*wfirst; 
        
        
        
        pctls = [w_p25; w_p75];
        pctls_idx = [find_w_p25; find_w_p75];
        target_pctl = [0.25; 0.75];
        avgw = [p_blm; p_abm];
        
        % Calculate the CH returns
        spillovers = zeros(2,1);
        for iter = 1:length(pctls)
      
            place_e_diff = zeros(N_n,N_n);
            place_e_diff_mass = zeros(N_n,1);
            for c = 1:length(Ss)    
                % Previous generation's neighborhood
                X_Prob = myDist0; 
                
                Probm = zeros(N_w,N_a);
                counter_neigh = setdiff(1:length(Ss), c);
                for k = 1:N_e
                    raw_Omega = unemp + (b0  + F_es0(:,:,c))  .* f_w(w0_1) ;
                    X_wprime0 = min(max(raw_Omega * e(k),w_min0),w_max0);    
                    [~, X_i_wprime0] = histc(X_wprime0,w0);
                    X_i_wprime0 = min(X_i_wprime0,N_w-1); 
                    weights = (X_wprime0-w0(X_i_wprime0))./(w0(X_i_wprime0+1)-w0(X_i_wprime0));
                    for j = 1:N_w
                        for i = 1:N_a % t+1 generation's ability
                            C1 = (1-weights(j,i))*X_Prob(j,i,c)*pi_e(k);
                            C2 =  weights(j,i)*X_Prob(j,i,c)*pi_e(k);
                            l_prime = X_i_wprime0(j,i);
                            for l=1:N_a % t+2 generation's ability
                                Probm(l_prime,l) = Probm(l_prime,l)+ C1*Pr(i,l);
                                Probm(l_prime+1,l) = Probm(l_prime+1,l)+ C2*Pr(i,l);
                            end
                         end
                     end
                end
               
                midx = pctls_idx(iter);
                mtarget = target_pctl(iter);
    
                full_wage_diff = (wprimeS);
                
                % Where their kids currently live
                idx = zeros(length(w0),1);
                idx(midx) = 1;
                myDist1_sub_1 = squeeze(sum(bsxfun(@times, Probm .* P_S .* idx, reshape(pi_theta, [1,1,1,2])),4));
                mfirst = sum(full_wage_diff .* squeeze(sum(myDist1_sub_1(:,:, counter_neigh),3)), 1:2) ./ sum(myDist1_sub_1(:,:, counter_neigh), 1:3);
                
                % Again, interpolating
                if cdf_w(midx) < mtarget
                    idx = zeros(length(w0),1);
                    idx(midx + 1) = 1;
                    myDist1_sub_2 = squeeze(sum(bsxfun(@times, Probm .* P_S .* idx, reshape(pi_theta, [1,1,1,2])),4));
                    wfirst = (mtarget - cdf_w(midx))/(cdf_w(midx+1) -  cdf_w(midx));
                    wlast = (cdf_w(midx+1) - mtarget)/(cdf_w(midx+1) -  cdf_w(midx));
                else
                    idx = zeros(length(w0),1);
                    idx(midx - 1) = 1;
                    myDist1_sub_2 = squeeze(sum(bsxfun(@times, Probm .* P_S .* idx, reshape(pi_theta, [1,1,1,2])),4));
                    wfirst = (cdf_w(midx) - mtarget)/(cdf_w(midx) -  cdf_w(midx-1));
                    wlast = (mtarget - cdf_w(midx-1))/(cdf_w(midx) -  cdf_w(midx-1));
                end
                mlast = sum(full_wage_diff .* squeeze(sum(myDist1_sub_2(:,:, counter_neigh),3)), 1:2) ./ sum(myDist1_sub_2(:,:, counter_neigh), 1:3);
                int_val = mfirst*wlast + mlast *wfirst;
                mass_int_val = wlast*myDist1_sub_1 + wfirst*myDist1_sub_2;
                
                % Average wage and wage mass
                place_e_diff(c, :) = int_val; 
                place_e_diff_mass(c) = sum(mass_int_val(:,:, counter_neigh), 1:3);
                place_e_diff(isnan(place_e_diff)) = 0;
 
            end
            
            % Std deviation for average wage of moving
            std_dev = sqrt(var(sum(place_e_diff .* place_e_diff_mass,1)./sum(place_e_diff_mass,1)));
            % Spillover
            spillovers(iter) = std_dev/avgw(iter); 
            
        end
        

        tmp3 = squeeze(sum(myDist1 .* w0, 1:2)./sum(myDist1 ,1:2));
        avg_inc_ratio = tmp3;

        
        dataset{t,1} = dissim;
        dataset{t,2} = gini;
        dataset{t,3} = avg_inc_ratio';
        dataset{t,4} = P_poor_S';
        dataset{t,6} = mobility_up;
        dataset{t,7} = N_size';
        dataset{t,8} = total_educ_return;
        dataset{t,9} = spillovers';
        dataset{t,10} = rs';
        dataset{t,11} = p_blm/p_abm;
        dataset{t,12} = cofv;
        dataset{t,13} = Ss';
        dataset{t,14} = P_rich_S' ./ N_size';
        dataset{t,15} = squeeze(share_C);
        dataset{t,16} = col_shares;
        dataset{t,17} = dissim_edu;
        dataset{t,18} =  sum((w0 < rs(3)) .* myDist1 ,1:4);
        dataset{t,19} = avg_ed;
        dataset{t,20} = avg_ability;
       
        myDist0 = myDist1;
        Prob = Prob_kids;
       
        F_es0 = F_es;
    end

end